import os
import argparse
import whisper
from moviepy.editor import VideoFileClip
from pydub import AudioSegment

# 确保 models 目录存在
os.makedirs("models", exist_ok=True)

def extract_audio_from_video(video_path, audio_path):
    """从视频文件中提取音频"""
    video = VideoFileClip(video_path)
    video.audio.write_audiofile(audio_path)
    video.close()

def convert_to_wav(input_path, output_path):
    """将音频转换为WAV格式"""
    audio = AudioSegment.from_file(input_path)
    audio.export(output_path, format="wav")

def extract_subtitles(file_path, output_format="srt", language="auto"):
    """提取字幕"""
    # 确定文件类型
    _, ext = os.path.splitext(file_path)
    ext = ext.lower()
    
    # 如果是视频文件，先提取音频
    temp_audio = "temp_audio.wav"
    if ext in ['.mp4', '.avi', '.mkv', '.mov']:
        temp_mp3 = "temp_audio.mp3"
        extract_audio_from_video(file_path, temp_mp3)
        convert_to_wav(temp_mp3, temp_audio)
        os.remove(temp_mp3)
    elif ext in ['.mp3', '.m4a', '.aac']:
        convert_to_wav(file_path, temp_audio)
    elif ext == '.wav':
        temp_audio = file_path
    else:
        raise ValueError(f"不支持的文件格式: {ext}")

    # 加载whisper模型
    print("正在加载Whisper模型...")
    model = whisper.load_model("small", download_root=os.path.join(os.getcwd(), "models"))
    
    # 转录音频
    print("正在转录音频...")
    
    # 设置转录参数
    transcribe_params = {
        'language': None if language == "auto" else language,
        'condition_on_previous_text': True,  # 考虑上下文
    }
    
    # 如果是中文，添加中文特定的提示
    if language in ["zh", "auto"]:
        transcribe_params.update({
            'initial_prompt': "请用简体中文输出，注意断句",
            'prompt': "以下是简体中文转录，注意断句："
        })
    
    result = model.transcribe(temp_audio, **transcribe_params)
    
    # 生成输出文件名
    output_file = os.path.splitext(file_path)[0] + f".{output_format}"
    
    # 保存字幕文件
    if output_format == "srt":
        with open(output_file, "w", encoding="utf-8") as f:
            subtitle_index = 1
            for segment in result["segments"]:
                text = segment["text"].strip()
                start_time = segment["start"]
                end_time = segment["end"]
                
                # 按空格分割文本（适用于英文和部分中文转录结果）
                words = text.split()
                lines = []
                current_line = ""
                
                # 按空格重组，控制每行长度
                for word in words:
                    # 如果当前行为空，直接添加单词
                    if not current_line:
                        current_line = word
                    # 如果添加这个词会导致行太长，或者词中包含句末标点，则换行
                    elif len(current_line) + len(word) + 1 > 20 or \
                         any(p in word for p in ['。', '！', '？', '.', '!', '?']):
                        lines.append(current_line)
                        current_line = word
                    # 否则，将词添加到当前行
                    else:
                        current_line += " " + word
                
                # 添加最后一行
                if current_line:
                    lines.append(current_line)
                
                # 如果没有成功分行，就直接用原文
                if not lines:
                    lines = [text]
                
                # 计算每行的时间
                duration = end_time - start_time
                time_per_line = duration / len(lines)
                
                # 写入字幕
                for i, line in enumerate(lines):
                    line_start = start_time + (i * time_per_line)
                    line_end = line_start + time_per_line
                    
                    f.write(f"{subtitle_index}\n")
                    f.write(f"{format_timestamp(line_start)} --> {format_timestamp(line_end)}\n")
                    f.write(f"{line.strip()}\n\n")
                    subtitle_index += 1
    else:
        with open(output_file, "w", encoding="utf-8") as f:
            f.write(result["text"])
    
    # 清理临时文件
    if temp_audio != file_path and os.path.exists(temp_audio):
        os.remove(temp_audio)
    
    print(f"字幕已保存到: {output_file}")

def format_timestamp(seconds):
    """将秒数转换为SRT时间戳格式"""
    hours = int(seconds // 3600)
    minutes = int((seconds % 3600) // 60)
    seconds = seconds % 60
    milliseconds = int((seconds - int(seconds)) * 1000)
    return f"{hours:02d}:{minutes:02d}:{int(seconds):02d},{milliseconds:03d}"

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="从视频或音频文件中提取字幕")
    parser.add_argument("file_path", help="视频或音频文件的路径")
    parser.add_argument("--format", choices=["srt", "txt"], default="srt", help="输出格式 (默认: srt)")
    parser.add_argument("--language", default="auto", help="音频语言 (默认: 自动检测)")
    
    args = parser.parse_args()
    extract_subtitles(args.file_path, args.format, args.language) 